package cn.jolyoulu.common.web.encrypt;

import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import cn.jolyoulu.common.core.utils.encryption.AESUtils;
import cn.jolyoulu.common.core.utils.encryption.Base64Utils;
import cn.jolyoulu.common.web.enums.GlobalExpType;
import cn.jolyoulu.common.web.excption.GlobalException;
import cn.jolyoulu.common.web.filter.MultiReadHttpServletRequestWrapper;
import cn.jolyoulu.common.web.filter.MultiReadHttpServletResponseWrapper;

import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @Author: JolyouLu
 * @Date: 2023/1/21 16:17
 * @Description 使用AES的加密解密过滤器
 */
@Slf4j
public class AesEncryptFilter implements Filter {
    /**
     * AES加密Key
     */
    private final String key;
    /**
     * body中timestamp超时时间
     * 所有请求都必须携带body，body中必须携带timestamp=当前时间戳参数，用于校验加密时间
     */
    private final Long bodyTimestampTimeOut;
    /**
     * 是否开启加密
     * false关闭加密过滤器，true开启
     */
    private final Boolean encryptEnable;
    /**
     * 跳过加密标记
     * 请求头中添加nonce与设定值相同时会跳过加密
     */
    private final String nonce;

    public AesEncryptFilter(String key, Boolean encryptEnable, String nonce, Long bodyTimestampTimeOut) {
        this.key = key;
        this.encryptEnable = encryptEnable;
        this.nonce = nonce;
        this.bodyTimestampTimeOut = bodyTimestampTimeOut;
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws ServletException, IOException {
        //判断是否是包装过的servlet
        if (servletRequest instanceof MultiReadHttpServletRequestWrapper
                && servletResponse instanceof MultiReadHttpServletResponseWrapper
                && isEncrypt(servletRequest)) {
            MultiReadHttpServletRequestWrapper requestWrapper = (MultiReadHttpServletRequestWrapper) servletRequest;
            MultiReadHttpServletResponseWrapper responseWrapper = (MultiReadHttpServletResponseWrapper) servletResponse;
            try {
                //对参数解密
                decryptParameter(requestWrapper, key);
                //对body解密
                decryptBody(requestWrapper, key);
                //timestamp时间校验
                check(requestWrapper, responseWrapper);
                //执行下一步过滤器
                filterChain.doFilter(servletRequest, servletResponse);
            } catch (Exception e) {
                log.error("AesEncryptFilter.doFilter 异常",e);
                servletResponse.getOutputStream().write("{\"code\":500,\"message\":\"加密/解密失败\"}".getBytes());
            } finally {
                //对body加密
                encryptBody(responseWrapper, key);
            }
        } else {
            filterChain.doFilter(servletRequest, servletResponse);
        }
    }

    /**
     * 对url中?后面的数据进行解密
     */
    private void decryptParameter(MultiReadHttpServletRequestWrapper requestWrapper, String key) {
        try {
            Map<String, String[]> parameterMap = requestWrapper.getParameterMap();
            if (Objects.nonNull(parameterMap) && !parameterMap.isEmpty()) {
                String ciphertext = parameterMap.keySet().iterator().next().trim();
                String decrypt = AESUtils.decrypt(ciphertext, key);
                for (String parameter : decrypt.split("&")) {
                    String[] kv = parameter.split("=");
                    requestWrapper.addParameter(kv[0], kv[1]);
                }
                requestWrapper.delParameter(ciphertext);
            }
        } catch (Exception e) {
            throw new GlobalException(GlobalExpType.ENCRYPT_ERROR);
        }
    }

    /**
     * 对body中的数据进行解密
     */
    private void decryptBody(MultiReadHttpServletRequestWrapper requestWrapper, String key) {
        try {
            byte[] deBase64 = Base64Utils.urlDecode(requestWrapper.getBytes());
            String decrypt = AESUtils.decrypt(deBase64, key);
            requestWrapper.setBytes(decrypt.getBytes());
            requestWrapper.setBody(decrypt);
        } catch (Exception e) {
            throw new GlobalException(GlobalExpType.ENCRYPT_ERROR);
        }
    }

    /**
     * 对body中的数据进行加密
     */
    private void encryptBody(MultiReadHttpServletResponseWrapper responseWrapper, String key) {
        try {
            String encrypt = AESUtils.encrypt(responseWrapper.toByteArray(), key);
            ByteArrayOutputStream byteArrayOutputStream = responseWrapper.getByteArrayOutputStream();
            //重置流
            byteArrayOutputStream.reset();
            byteArrayOutputStream.write(encrypt.getBytes());
        } catch (Exception e) {
            throw new GlobalException(GlobalExpType.ENCRYPT_ERROR);
        }
    }

    /**
     * 初步校验密文
     */
    private void check(MultiReadHttpServletRequestWrapper requestWrapper, MultiReadHttpServletResponseWrapper responseWrapper) {
        String method = requestWrapper.getMethod();
        //判断请求类型
        Long timestamp = 0L;
        if (method.equals("GET")) {
            //校验Parameter
            Map<String, String[]> parameterMap = requestWrapper.getParameterMap();
            String timestampParam = parameterMap.get("timestamp")[0];
            if (StringUtils.isNotBlank(timestampParam)) {
                timestamp = Long.valueOf(timestampParam);
            }
        } else {
            //校验Body
            List<String> contentTypes = Arrays.asList(requestWrapper.getContentType().split(";"));
            if (contentTypes.contains("application/json")) {
                JSONObject jsonObject = JSONObject.parseObject(requestWrapper.getBody());
                timestamp = jsonObject.getLong("timestamp");
            }
        }
        if (timestamp <= 0) {
            throw new GlobalException(GlobalExpType.ENCRYPT_ERROR);
        }
        //60秒前的数据统一不处理，大于当前时间的数据也不处理
        if (Math.abs(System.currentTimeMillis() - timestamp) > bodyTimestampTimeOut) {
            throw new GlobalException(GlobalExpType.ENCRYPT_ERROR);
        }
    }

    /**
     * 是否需要加密
     *
     * @return true开启加密, false跳过加密
     */
    private boolean isEncrypt(ServletRequest servletRequest) {
        //判断是否全局关闭加密
        if (this.encryptEnable.equals(false)){
            return false;
        }
        //判断请求头的nonce跳过加密
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        String nonce = request.getHeader("nonce");
        if (StringUtils.isNoneBlank(nonce) && nonce.equals(this.nonce)) {
            return false;
        }
        return true;
    }

}
